import streamlit as st
import cv2
import numpy as np
from ultralytics import YOLO
from PIL import Image
import tempfile

# 页面配置
st.set_page_config(
    page_title="目标检测系统",
    page_icon="🎯",
    layout="wide"
)

# 标题
st.title("🎯 实时目标检测系统")

# 初始化模型
@st.cache_resource
def load_model():
    return YOLO('yolov8n.pt')

model = load_model()

# 侧边栏
detection_mode = st.sidebar.selectbox(
    "选择检测模式",
    ["图片检测", "视频检测", "实时检测"]
)

confidence = st.sidebar.slider(
    "置信度阈值",
    min_value=0.0,
    max_value=1.0,
    value=0.5
)

def process_image(image, model, conf):
    results = model.predict(image, conf=conf)[0]
    return results

def draw_boxes(image, results):
    for result in results.boxes.data.tolist():
        x1, y1, x2, y2, score, class_id = result
        
        if score > confidence:
            cv2.rectangle(image, 
                         (int(x1), int(y1)), 
                         (int(x2), int(y2)), 
                         (0, 255, 0), 2)
            
            label = f"{results.names[int(class_id)]} {score:.2f}"
            cv2.putText(image, label, (int(x1), int(y1)-10), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
    
    return image

if detection_mode == "图片检测":
    uploaded_file = st.file_uploader("上传图片", type=['jpg', 'jpeg', 'png'])
    
    if uploaded_file is not None:
        image = Image.open(uploaded_file)
        image_np = np.array(image)
        
        col1, col2 = st.columns(2)
        with col1:
            st.subheader("原始图片")
            st.image(image)
        
        results = process_image(image_np, model, confidence)
        output_image = draw_boxes(image_np.copy(), results)
        
        with col2:
            st.subheader("检测结果")
            st.image(output_image)

elif detection_mode == "视频检测":
    uploaded_file = st.file_uploader("上传视频", type=['mp4', 'avi', 'mov'])
    
    if uploaded_file is not None:
        temp_file = tempfile.NamedTemporaryFile(delete=False)
        temp_file.write(uploaded_file.read())
        
        cap = cv2.VideoCapture(temp_file.name)
        stframe = st.empty()
        
        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break
                
            results = process_image(frame, model, confidence)
            output_frame = draw_boxes(frame.copy(), results)
            stframe.image(output_frame, channels="BGR")
            
        cap.release()

elif detection_mode == "实时检测":
    run = st.checkbox('开启摄像头')
    FRAME_WINDOW = st.image([])
    camera = cv2.VideoCapture(0)

    while run:
        _, frame = camera.read()
        if frame is not None:
            results = process_image(frame, model, confidence)
            output_frame = draw_boxes(frame.copy(), results)
            FRAME_WINDOW.image(output_frame, channels="BGR")
    else:
        camera.release()